from .pd_resnet import pdresnet18
import torch.nn as nn
import numpy as np
from collections import OrderedDict
from .layers import *
import torchvision.models as models


class ResnetEncoder_pcn(nn.Module):
    """Pytorch module for a resnet encoder
    """
    def __init__(self,input_channel=3):
        super(ResnetEncoder_pcn, self).__init__()
        self.encoder=pdresnet18(input_channel=input_channel)
    def forward(self, input_image,mask):
        self.features = []
        self.masks=[]
        x=input_image
        x,mask = self.encoder.conv1(x,mask)
        x = self.encoder.bn1(x)
        x=self.encoder.relu(x)
        self.features.append(x)
        self.masks.append(mask)
        x,mask=self.encoder.layer1[0](self.encoder.maxpool(x),self.encoder.maxpool(mask))
        x,mask=self.encoder.layer1[1](x,mask)
        self.features.append(x)
        self.masks.append(mask)
        x,mask=self.encoder.layer2[0](x,mask)
        x,mask=self.encoder.layer2[1](x,mask)
        self.features.append(x)
        self.masks.append(mask)
        x,mask=self.encoder.layer3[0](x,mask)
        x,mask=self.encoder.layer3[1](x,mask)
        self.features.append(x)
        self.masks.append(mask)
        x,mask=self.encoder.layer4[0](x,mask)
        x,mask=self.encoder.layer4[1](x,mask)
        self.features.append(x)
        self.masks.append(mask)
        return self.features, self.masks

class Decoder_PCN(nn.Module):
    def __init__(self, num_output_channels=3):
        super(Decoder_PCN, self).__init__()
        self.num_ch_enc = np.array([64, 64, 128, 256, 512])
        self.num_output_channels = num_output_channels
        self.use_skips = True
        # self.upsample_mode = 'nearest'
        self.scales = [0]
        self.num_ch_dec = np.array([16, 32, 64, 128, 256])

        # decoder
        self.convs = OrderedDict()
        for i in range(4, -1, -1):
            # upconv_0
            num_ch_in = self.num_ch_enc[-1] if i == 4 else self.num_ch_dec[i + 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 0)] = ConvBlock_PCN(num_ch_in, num_ch_out)
            # upconv_1
            num_ch_in = self.num_ch_dec[i]
            if self.use_skips and i > 0:
                num_ch_in += self.num_ch_enc[i - 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 1)] = ConvBlock_PCN(num_ch_in, num_ch_out)
        for s in self.scales:
            self.convs[("dispconv", s)] = Conv3x3_PCN(self.num_ch_dec[s], self.num_output_channels)

        self.decoder = nn.ModuleList(list(self.convs.values()))
        
    def forward(self, input_features,masks):
        self.outputs = []
        
        x = input_features[-1]
        mask=masks[-1]
        for i in range(4, -1, -1):
            x,mask= self.convs[("upconv", i, 0)](x,mask)
            x = [upsample(x)]
            mask = upsample(mask)
            if self.use_skips and i > 0:
                x += [input_features[i - 1]]
            x = torch.cat(x, 1)
            x,mask= self.convs[("upconv", i, 1)](x,mask)
            if i in self.scales:
                x,mask=self.convs[("dispconv", i)](x,mask)
                self.outputs.append(x)
        return  self.outputs[-1]

class Unet_pcn(nn.Module):
    def __init__(self,cfg):
        super().__init__()
        self.cfg = cfg
       
        self.encoder=ResnetEncoder_pcn(input_channel=cfg.frm_input)
        self.decoder=Decoder_PCN(num_output_channels=cfg.frm_output)

    def forward(self, img,mask):
        feature,masks=self.encoder(img,mask)

        output=self.decoder(feature,masks)+img

        return  output